﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;

namespace FileArchive
{
    class ArchiveNode
    {
        public string Name { get; set; }
        internal virtual void SaveInternal(BinaryWriter writer, BinaryWriter bucketWriter, ICryptoTransform encryptor, ref long offset)
        {
            var nameBytes = Encoding.Unicode.GetBytes(Name);
            var transformedNameBytes = encryptor.TransformFinalBlock(nameBytes, 0, nameBytes.Length);
            writer.Write(transformedNameBytes.Length);
            writer.Write(transformedNameBytes);
        }
        internal virtual void LoadInternal(BinaryReader reader, ICryptoTransform decryptor)
        {
            int len = reader.ReadInt32();
            var encryptedBytes = reader.ReadBytes(len);
            var nameBytes = decryptor.TransformFinalBlock(encryptedBytes, 0, encryptedBytes.Length);
            Name = Encoding.Unicode.GetString(nameBytes);
        }
        public FolderNode Parent { get; set; }
    }
    class FileNode : ArchiveNode
    {
        public string OriginalFileName;
        long fileOffset;
        int size;
        internal override void SaveInternal(BinaryWriter writer, BinaryWriter bucketWriter, ICryptoTransform encryptor, ref long offset)
        {
            base.SaveInternal(writer, bucketWriter, encryptor, ref offset);
            writer.Write(offset);
            
            var bytes = File.ReadAllBytes(OriginalFileName);
            var transformedBytes = encryptor.TransformFinalBlock(bytes, 0, bytes.Length);
            bucketWriter.Write(transformedBytes);
            size = transformedBytes.Length;
            writer.Write(size);
            offset += transformedBytes.Length;

        }
        internal override void LoadInternal(BinaryReader reader, ICryptoTransform decryptor)
        {
            base.LoadInternal(reader, decryptor);
            fileOffset = reader.ReadInt64();
            size = reader.ReadInt32();
        }
        internal byte[] GetContent(Stream s, ICryptoTransform decryptor, long bucketOffset)
        {
            s.Seek(bucketOffset+fileOffset, SeekOrigin.Begin);
            byte[] buffer = new byte[size];
            s.Read(buffer, 0, size);
            var nBytes = decryptor.TransformFinalBlock(buffer, 0, buffer.Length);
            return nBytes;
        }
    }
    class FolderNode : ArchiveNode
    {
        public List<ArchiveNode> Children { get; private set; }
        public FolderNode()
        {
            Children = new List<ArchiveNode>();
        }
        internal override void SaveInternal(BinaryWriter writer, BinaryWriter bucketWriter, ICryptoTransform encryptor, ref long offset)
        {
            base.SaveInternal(writer, bucketWriter, encryptor, ref offset);
            writer.Write(Children.Count);
            foreach (var child in Children)
            {
                if (child is FolderNode)
                    writer.Write(true);
                else
                    writer.Write(false);
                child.SaveInternal(writer, bucketWriter, encryptor, ref offset);
            }
        }
        internal override void LoadInternal(BinaryReader reader, ICryptoTransform decryptor)
        {
            base.LoadInternal(reader, decryptor);
            int count = reader.ReadInt32();
            for (int i = 0; i<count; i++)
            {
                ArchiveNode child;
                if (reader.ReadBoolean())
                {
                    child = new FolderNode();
                }
                else
                {
                    child = new FileNode();
                }
                child.Parent = this;
                child.LoadInternal(reader, decryptor);
                Children.Add(child);
            }
        }
        public void Save(BinaryWriter writer, ICryptoTransform encryptor)
        {
            using (var bucketWriter = new BinaryWriter(new FileStream("~temp.bin", FileMode.Create)))
            {
                long count = 0;
                SaveInternal(writer, bucketWriter, encryptor, ref count);
            }
            using (var bucketStream = new FileStream("~temp.bin", FileMode.Open))
            {
                writer.Write(writer.BaseStream.Length + 8);
                bucketStream.CopyTo(writer.BaseStream);
            }
            File.Delete("~temp.bin");
        }

        internal long Load(BinaryReader reader, ICryptoTransform decryptor)
        {
            LoadInternal(reader, decryptor);
            long bucketOffset = reader.ReadInt64();
            return bucketOffset;
        }

    }
    
    class Archive
    {
        System.Security.Cryptography.RijndaelManaged provider = new System.Security.Cryptography.RijndaelManaged();
        FolderNode rootNode = new FolderNode();
        ICryptoTransform encryptor, decryptor;
        long bucketOffset;
        public Archive(byte[] key)
        {
            provider.Key = key;
            encryptor = provider.CreateEncryptor();
            decryptor = provider.CreateDecryptor();
        }

        public void Save(BinaryWriter writer)
        {
            writer.Write(provider.IV.Length);
            writer.Write(provider.IV);
            byte[] bytes = BitConverter.GetBytes(19890911);
            var transBytes = encryptor.TransformFinalBlock(bytes, 0, bytes.Length);
            writer.Write(transBytes.Length);
            writer.Write(transBytes);
            rootNode.Save(writer, encryptor);
        }
        Stream stream = null;
        public Archive(BinaryReader reader, byte[] key)
        {
            var len = reader.ReadInt32();
            provider.IV = reader.ReadBytes(len);
            provider.Key = key;
            encryptor = provider.CreateEncryptor();
            decryptor = provider.CreateDecryptor();
            try
            {
                len = reader.ReadInt32();
                var transBytes = reader.ReadBytes(len);
                var bytes = decryptor.TransformFinalBlock(transBytes, 0, len);
                var trueBytes = BitConverter.GetBytes(19890911);
                if (bytes.Length != trueBytes.Length)
                    throw new InvalidOperationException();
                for (int i = 0; i < bytes.Length; i++)
                    if (bytes[i] != trueBytes[i])
                        throw new InvalidOperationException();
                stream = reader.BaseStream;
            }
            catch (Exception)
            {
                throw new InvalidOperationException();
            }
            bucketOffset = rootNode.Load(reader, decryptor);
        }

        public Stream LoadFile(FileNode node)
        {
            lock (stream)
            {
                var bytes = node.GetContent(stream, decryptor, bucketOffset);
                return new MemoryStream(bytes);
            }
        }

        private static void BuildArchiveFromDirectory(string path, FolderNode node)
        {
            var dirs =  Directory.EnumerateDirectories(path);
            foreach (var dir in dirs)
            {
                var nNode = new FolderNode();
                nNode.Name = Path.GetFileName(dir);
                node.Children.Add(nNode);
                nNode.Parent = node;
                BuildArchiveFromDirectory(dir, nNode);
            }
            foreach (var file in Directory.EnumerateFiles(path))
            {
                var nNode = new FileNode();
                nNode.Name = Path.GetFileName(file);
                nNode.OriginalFileName = file;
                nNode.Parent = node;
                node.Children.Add(nNode);
            }
        }

        public static Archive FromDirectory(string path, string password)
        {
            Archive result = new Archive(GenerateKeyFromPassword(password));
            result.rootNode.Name = Path.GetFileName(path);
            BuildArchiveFromDirectory(path, result.rootNode);
            return result;
        }

        public static Archive FromFile(string fileName, string password)
        {
            Archive archive;
            FileStream fs = new FileStream(fileName, FileMode.Open, FileAccess.Read, FileShare.Read);
            {
                fs.Seek(-8, SeekOrigin.End);
                var reader = new BinaryReader(fs);
                {
                    long offset = reader.ReadInt64();
                    fs.Seek(offset, SeekOrigin.Begin);
                    archive = new Archive(reader, GenerateKeyFromPassword(password));
                }
            }
            return archive;
        }
        static byte[] GenerateKeyFromPassword(string password)
        {
            byte[] rs = new byte[32];
            var bytes = Encoding.Unicode.GetBytes(password);
            for (int i = 0; i < rs.Length; i++)
                rs[i] = bytes[i % bytes.Length];
            return rs;
        }
        public override string ToString()
        {
            StringBuilder sb = new StringBuilder();
            sb.AppendLine(rootNode.Name);
            foreach (var child in rootNode.Children)
            {
                sb.Append("  ");
                sb.AppendLine(child.Name);
            }
            return sb.ToString();
        }

        public FolderNode Root { get { return rootNode; } }


        internal void Dispose()
        {
            if (stream != null)
                stream.Dispose();
        }
    }
}
